package com.xl.ssh.service.impl;

import com.alibaba.fastjson.JSON;
import com.jcraft.jsch.*;
import com.xl.ssh.constants.Constants;
import com.xl.ssh.domain.SSHConnectDO;
import com.xl.ssh.domain.WebSSHData;
import com.xl.ssh.service.WebSSHService;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import javax.annotation.Resource;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;

@Component
public class WebSSHServiceImpl implements WebSSHService {

    //存储用于session，session可用于发送消息.键是用户连接时候生成的uuid，值是SSHConnectDo对象，里面存储这各种ssh通道信息
    private static Map<String, SSHConnectDO> userConnetions = new ConcurrentHashMap<>();

    @Resource//根据名字注入
    private ThreadPoolTaskExecutor sshTaskThreadPool;


    /**
     * 存储用户的连接信息，放到map集合里面方便后期使用
     *
     * @param session
     */
    @Override
    public void initConnection(WebSocketSession session) {
        JSch jSch = new JSch();
        SSHConnectDO sshConnectDO = new SSHConnectDO();
        sshConnectDO.setjSch(jSch);
        sshConnectDO.setWebSocketSession(session);
        String uuid = session.getAttributes().get(Constants.USER_UUID_KEY).toString();
        userConnetions.put(uuid, sshConnectDO);
    }

    /**
     * 接收到前端发送过来的shell命令，进行执行
     *
     * @param buffer
     * @param session
     */
    @Override
    public void recvHandle(String buffer, WebSocketSession session) {
        WebSSHData webSSHData = null;
        //解析前端发送过来的数据
        try {
            webSSHData = JSON.parseObject(buffer, WebSSHData.class);
        } catch (Exception e) {
            e.printStackTrace();
        }
        System.out.println("收到的消息：【"+webSSHData.toString()+"】");
        String user_uuid = session.getAttributes().get(Constants.USER_UUID_KEY).toString();
//        System.out.println("uuid:" + user_uuid);
//        System.out.println("username:" + session.getAttributes().get(Constants.USER_INFO_KEY));
//        System.out.println("ssh_command:" + webSSHData.getCommand());

        //如果是连接操作进行ssh连接
        if (Constants.WEBSSH_OPERATE_CONNECT.equals(webSSHData.getOperate())) {
            //通过uuid找到socket连接时初始化的时候，SSHConnectDO，里面包含了链接建立时候的session，JSch ，和channel对象（用于发送shell，此时为null）
            SSHConnectDO sshConnectDO = userConnetions.get(user_uuid);
            //开始建立ssh连接，因为ssh连接是阻塞式的，要使用子线程，一个子线程表示一个ssh信道
            WebSSHData finalWebSSHData = webSSHData;
            sshTaskThreadPool.execute(()->{
                try {
                    createSSHConnection(sshConnectDO, finalWebSSHData, session);
                } catch (JSchException  | IOException e) {
                    System.err.println("ssh连接异常");
                    System.err.println(e);
                    //错误信息发送给前端
                    try {
                        sendMessage(session,"ssh连接异常\n".getBytes());
                    } catch (IOException ex) {
                        ex.printStackTrace();
                    }
                }

            });
        } else if (Constants.WEBSSH_OPERATE_COMMAND.equals(webSSHData.getOperate())) {
            //如果是执行命令的操作 ,本质还是把命令发送给终端
            String cmmand = webSSHData.getCommand();
            SSHConnectDO sshConnectDO = userConnetions.get(user_uuid);
            if ( sshConnectDO != null){
                try {
                    ChannelShell channel = (ChannelShell)sshConnectDO.getChannel();
                    channel.setPtySize(webSSHData.getCols() , webSSHData.getRows(),  640 , 480);
                    sendToChannel(channel,cmmand);
                } catch (IOException e) {
                    System.err.println("异常信息："+e.getMessage());
                    try {
                        sendMessage(session, "ssh连接异常\n".getBytes());
                    } catch (IOException ioException) {
                    }
                    e.printStackTrace();
                }
            }
        }else {
            System.err.println("不支持的操作");
            close(session);
        }

    }


    /**
     * 方法里面肯定有一个是阻塞式的
     *  InputStream inputStream = channelShell.getInputStream(); 这应该是阻塞式的，一直就卡在这里了，所以要使用子线程处理
     */
    private void createSSHConnection(SSHConnectDO sshConnectDO, WebSSHData webSSHData, WebSocketSession webSocketSession) throws JSchException, IOException {
        Session session = null ;
        Properties config = new Properties();
        config.put("StrictHostKeyChecking","no");
        //获取ssh会话
        session = sshConnectDO.getjSch().getSession(webSSHData.getUsername(), webSSHData.getHost(), webSSHData.getPort());
        session.setConfig(config);
        //设置密码
        session.setPassword(webSSHData.getPassword());
        //连接通道设置超时时间为20s
        session.connect(20*1000);
        //开启shell通道
        ChannelShell channelShell = (ChannelShell)session.openChannel("shell");
        //设置连接通道的超时时间3s
        channelShell.connect(3*1000);
        //设置shell窗口的大小?
        channelShell.setPtySize(webSSHData.getCols(),webSSHData.getRows(),640,480);
        //设置channel
        sshConnectDO.setChannel(channelShell);
        //转发消息-终端执行command
        sendToChannel(channelShell , "\r");
        //获取终端返回的数据
        InputStream inputStream = channelShell.getInputStream();
        //将终端返回的数据，发送给前端
        try {
            byte[]  buffer = new byte[1024];
            int len = 0  ;
            while (( len = inputStream.read(buffer))  != -1){
                sendMessage(webSocketSession , Arrays.copyOfRange(buffer,0,len));
            }
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            //关闭会话
            session.disconnect();
            channelShell.disconnect();
            if (inputStream != null){
                inputStream.close();
            }
        }

    }

    private void sendToChannel(ChannelShell channelShell, String command) throws IOException {
        if (channelShell != null){
            OutputStream outputStream = channelShell.getOutputStream();
            outputStream.write(command.getBytes());
            outputStream.flush();
        }
    }

    @Override
    public void sendMessage(WebSocketSession session, byte[] buffer) throws IOException {
        session.sendMessage(new TextMessage(buffer));
    }

    //ws连接关闭后，断开 shell 的session 和channel通道
    @Override
    public void close(WebSocketSession session) {
        String user_uuid = session.getAttributes().get(Constants.USER_UUID_KEY).toString();
        SSHConnectDO sshConnectDO = userConnetions.get(user_uuid);
        if (sshConnectDO != null){
            if (sshConnectDO.getChannel() != null ){
                sshConnectDO.getChannel().disconnect();
            }
            userConnetions.remove(user_uuid);
        }
    }
}

